package comlib

import chisel3._
import chisel3.util._

object util{

  /** Onehot控制的MUX选择
   * @param in 被选择的数据数据列表
   * @param cgnt Onehot编码的选择信号. Note: 若cgnt==0，则该函数输出数据也为"0"
   * @return 选择的数据
   * @group XB
   */
  def muxO[B<:Data](in:Vec[B], cgnt:Vec[Bool]):B = {
    val ins = in.map(_.asUInt())
    val s = ins(0).getWidth
    cgnt
      .map {case b => VecInit(Seq.fill(s)(b))}
      .map {case b => b.asUInt()}
      .zip (ins)
      .map {case (s,c) => s&c}
      .reduce (_|_)
      .asTypeOf( in(0) )
  }

    /** 带2Mux1选择器的寄存器
      * 
      * @param sel Mux选择信号
      * @param a sel==true.B时，选择寄存的信号
      * @param b sel==false.B时，选择寄存的信号
      * @return Mux选择器后的寄存器
      */
    def RegMux[T<:Data]( sel:Bool, a:T, b:T ) = RegNext( Mux(sel, a, b) )

    /** 带使能信号的寄存器。不是能时输出为0.U
      *
      * @param en 使能信号。true.B使能
      * @param next 寄存器Q端信号
      * @return 寄存器
      */
    def RegDisable0[T<:Data]( en:Bool, next:T ) = RegMux( en, next, 0.U.asInstanceOf[T] )
    
    /** 带使能信号的寄存器。不是能时输出为0xF..F
      *
      * @param en 使能信号。true.B使能
      * @param next 寄存器Q端信号
      * @return 寄存器
      */
    def RegDisable1[T<:Data]( en:Bool, next:T ) = RegMux( en, next, Seq.fill(next.getWidth)("1").mkString("b","","").U.asInstanceOf[T] )

    /** 根据lock参数选择next是否寄存器输出
      *
      * @param lock Register output flag
      * @param next output value
      * @return Register if lock==true, or Wire if lcok==false
      */
    def Connect[T<:Data]( lock:Boolean, next:T ) = {
        val ret = if( lock ) Reg(chiselTypeOf(next)) else Wire(chiselTypeOf(next))
        ret := next
        ret
    }

    /** Output `next` value using Register when `lock` is true, or Wire when `lock` is false
      * @param lock Register output flag parameter
      * @param init Initial value for register
      * @param next output signal
      */
    def ConnectInit[T<:Data]( lock:Boolean, init:T, next:T ) = {
        val ret = if( lock ) RegInit(init) else WireInit(init)
        ret := next
        ret
    }

    /** 计算"1"的个数
     * 
     * 采用二分法并行计算，每个节点采用进位加法
     *
     * @tparam T 输入数据类型。必须为chisel类型
     * @param d 输入数据
     * @param lockStage 二分法参数。参考[[reduce]]
     * @param stage 二分法参数。参考[[reduce]]
     * @return 数据`d`中的"1"的个数
     */
    def oneNumber[T<:Data]( d:T, lockStage:Int=0, stage:Int=1 ):UInt = {
      val s = d.asUInt().asBools() map {case x => Mux(x, 1.U, 0.U)}
      reduce.binaryTree( (x:UInt,y:UInt)=>x+&y, s, lockStage, stage )
    }

    def zeroNumber[T<:Data]( d:T, lockStage:Int=0, stage:Int=1 ):UInt = {
      val s = d.asUInt().asBools() map {case x => Mux(x, 0.U, 1.U)}
      reduce.binaryTree( (x:UInt,y:UInt)=>x+&y, s, lockStage, stage )
    }

    /** 按bit转置
     */
    def transpose( in:Vec[UInt] ) : Vec[UInt] = {
      VecInit.tabulate(in(0).getWidth){case c => Cat( for( r <- in.size-1 to 0 by -1 ) yield in(r)(c) )}
    }

    def isOnehot(in:UInt) = in&(in-1.U)===0.U

    //===========================================================================================
    // fillSlice
    class FillMode
    object FillNone extends FillMode
    object Fill0 extends FillMode 
    object Fill1 extends FillMode
    object FillDest extends FillMode
    object FillExt extends FillMode

    /** 数据部分Bit赋值
     * @param d 目的数据
     * @param s 源数据
     * @param st 目的数据赋值起始位置(include)
     * @param ed 目的数据赋值结束位置(include)
     * @return 赋值完成后的A类型数据
     * @note 当源数据的bit宽度小于需要赋值的宽度(st-ed+1)时，按LSB对齐赋值，高位补"0"
     * @example 
     {{{
       fillSlice("b10000_00000000_001".U, "b1010_1010".U, 10, 3)  //ans is "b10000_10101010_001".U
     }}}
     */
    def fillSlice[A<:Data,B<:Data](d:A, s:B, st:Int, ed:Int, fill:FillMode=Fill0):A = {
      val ss = s.asUInt().asBools()
      val ds = d.asUInt().asBools()
      require( st < ds.size && ed <= st && ed >= 0 )
      Cat(
        for(i <- ds.size-1 to 0 by -1) yield
          if( i > st || i < ed ) ds(i)
          else if( i - ed >= ss.size ) fill match{
            case Fill0 => false.B 
            case Fill1 => true.B
            case FillDest => ds(i)
            case FillNone => throw new ArithmeticException("You are not eligible")
          }
          else ss(i-ed) 
      ).asTypeOf(d)
    }

    def fillSlice0[A<:Data,B<:Data](d:A, s:B, st:Int, ed:Int):A = fillSlice(d,s,st,ed,Fill0)
    def fillSlice0[A<:Data,B<:Data](d:A, s:B, st:Int):A = fillSlice(d,s,st,0,Fill0)
    def fillSlice1[A<:Data,B<:Data](d:A, s:B, st:Int, ed:Int):A = fillSlice(d,s,st,ed,Fill1)
    def fillSlice1[A<:Data,B<:Data](d:A, s:B, st:Int):A = fillSlice(d,s,st,0,Fill1)

    //================================================================================================
    // Group
    /** 把数据s按照gen为模板拆分到一个Vec中
     * @param gen 输出模板数据
     * @param s 源数据
     * @param fill 填充行为。目前支持
     *  - Fill0: 用"0"填充不足位
     *  - Fill1: 用"1"填充不足位
     *  - FillExt: 用s的高位填充不足位
     * @return 返回目标Vec
     */
    def group[A<:Data,B<:Data](gen:A, s:B, fill:FillMode):Vec[A] = {
      val su = s.asUInt()
      val dw = gen.asUInt().getWidth
      val fillVal = fill match {
        case Fill0 => 0.U(1.W)
        case Fill1 => 1.U(1.W)
        case FillExt => su(su.getWidth-1)
        case _ => 0.U(1.W)
      }

      val ss = su.getWidth % dw match {
        case 0 => su
        case r => Cat(su, Fill(dw-r, fillVal ))
      }

      VecInit(ss.asBools().grouped(dw).map( VecInit(_).asUInt().asTypeOf(chiselTypeOf(gen)) ).toSeq)
    }

    def groupMux[D<:Data](sel:Vec[Bool], in1:D, in2:D):D = {
      val width = in1.getWidth 
      val splitNum = sel.size
      val fullDataSize = scala.math.ceil(width/splitNum).toInt 
      val rmnDataSize = width - fullDataSize * (splitNum-1)
      val buf1 = in1.asUInt 
      val buf2 = in2.asUInt 

      val fullData = (splitNum-1 to 1 by -1).map{ case s => 
        val mux_lsb = fullDataSize * (s-2) + rmnDataSize
        val mux_msb = mux_lsb + fullDataSize - 1
        Mux( sel(s), buf1(mux_msb, mux_lsb), buf2(mux_msb, mux_lsb) )
      }

      val rmnData = Mux( sel(0), buf1(rmnDataSize-1, 0), buf2(rmnDataSize-1, 0) )

      Cat( fullData :+ rmnData ).asTypeOf(in1)
    }

    def groupSelect[D<:Data](sel:Vec[UInt], in:Vec[D]):D = {
      val width = in(0).getWidth 
      val splitNum = sel.size
      val fullDataSize = scala.math.ceil(width/splitNum).toInt 
      val rmnDataSize = width - fullDataSize * (splitNum-1)
      val bufs = VecInit(in.map(_.asUInt))

      val fullData = (splitNum-1 to 1 by -1).map{ case s => 
        val mux_lsb = fullDataSize * (s-2) + rmnDataSize
        val mux_msb = mux_lsb + fullDataSize - 1
        bufs( sel(s) )( mux_msb, mux_lsb )
      }

      val rmnData = bufs( sel(0) )( rmnDataSize-1, 0 )

      Cat( fullData :+ rmnData ).asTypeOf(in(0))
    }


    implicit def IndexedSeq_To_ChiselSeq[T<:Data]( s:Seq[T] ) = new ChiselSeq(s)
    implicit def Data_To_ComlibData[T<:Data]( s:T ) = new ComlibData(s)

  /** 循环左移 
   * @param in 左移操作数
   * @param ofst 左移位数
   * @return 循环左移后的结果
   * @group ALU
   */
  def rrShiftL(in:UInt, ofst:Int):UInt = {
    val in_inner = in.asBools()
    val size = in.getWidth
    val res = (0 until size).map{ case i => in((i+size-ofst)%size) }
    VecInit(res).asUInt
  }

  /** 循环右移 
   * @param in 右移操作数
   * @param ofst 右移位数
   * @return 循环右移后的结果
   * @group ALU
   */
  def rrShiftR(in:UInt, ofst:Int):UInt = {
    val in_inner = in.asBools()
    val size = in.getWidth
    val res = (0 until size).map{ case i => in((i+ofst)%size) }
    VecInit(res).asUInt
  }

  /** Onehot编码转为化温度计编码
   *
   * @param in onehot编码输入信号
   * @return 温度计编码
   * @group Format
   */
  def OH2Temp(in:UInt) = {
    val s = in.getWidth
    val low = VecInit((0 until s-1).map{ case i => in(s-1,i+1).asBools().reduce(_|_)}).asUInt
    Cat(0.U(1.W), low)
  }

  /** 二进制编码到温度计编码转化
   *
  * 转化规则：
  *  - 0 => "b0000_0000"
  *  - 1 => "b0000_0001"
  *  - 2 => "b0000_0011"
  *  - 3 => "b0000_0111"
  *  - 4 => "b0000_1111"
  *  - 5 => "b0001_1111"
  *  - 6 => "b0011_1111"
  *  - 7 => "b0111_1111"
  *
  * @param in 输入的二进制编码
  * @param maxPos 最大表示数
  * @return 温度计编码格式
  * @group Format
   */
  def bin2Temp( in:UInt, maxPos:Int ) : Vec[Bool] = {
    val zero = 0.U(maxPos.W)
    val one = ~((~0.U(maxPos.W)) << in)
    VecInit((zero | one.asUInt()).asBools())
  }

  /** 转化onehot编码为binary编码
   * @param in onehot编码数据。该数据由Vec[Bool]的数据结构来表示
   * @param e 若in为全"0"值时，默认的输出数据
   * @return 返回包含2个元素的Tuple数据结构
   *   - Bool: in中是否有"true.B"
   *   - UInt: 转化后的binary编码值
   * @group util
   */
  def OH2Bin(in:Vec[Bool], e:UInt):(Bool,UInt) = {
    val has = in.reduce(_|_)
    (has, Mux(has, OHToUInt(in.asUInt), e))
  }

  /** 按bit取反 */
  def reverse( in:UInt ) = VecInit(in.asBools.reverse).asUInt

}

